﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace Assignment4
{
    class ID3Tree
    {
        public static Dictionary<string, List<Example>> splitByAttribute(List<Example> examples, string attributeType)
        {
            Dictionary<string, List<Example>> sets = new Dictionary<string, List<Example>>();

            //# We create a set of lists, so we can access each list
            //# by the attribute value
            List<string> values;
            if (Example.AttributeCollection.TryGetValue(attributeType, out values))
            {
                foreach (string value in values)
                {
                    sets.Add(value, new List<Example>());
                }
            }
            else
            {
                Console.WriteLine("There is no attribute named {0} in dictionary", attributeType);
                Console.ReadLine();
                Environment.Exit(0);
            }

            //# Loop through each example
            foreach (Example example in examples)
            {
                //# Assign it to the right set
                string exampleValue = example.getValue(attributeType);
                List<Example> exampleMatched = sets[exampleValue];
                exampleMatched.Add(example);
                sets[exampleValue] = exampleMatched;
            }
            //#Return the sets
            return sets;
        }

        public static float entropy(List<Example> examples)
        {
            //# Get the number of examples
            int exampleCount = examples.Count;

            //# Check if we only have one: in that case entropy is 0
            if (exampleCount < 2)
                return 0;
            //# Otherwise we need to keep a tally of how many of each
            //# different kind of action we have
            Dictionary<string, int> actionTallies = new Dictionary<string, int>();

            //# Go through each example
            foreach (Example example in examples)
            {
                try
                {
                    //# Increment the appropriate tally
                    actionTallies[example.Action]++;
                }
                catch (KeyNotFoundException)
                {
                    actionTallies.Add(example.Action, 1);
                }
            }

            //# We now have the counts for each action in the set
            int actionCount = actionTallies.Keys.Count;

            //# If we have only one action then we have zero entropy
            if (actionCount < 2) return 0;

            //# Start with zero entropy
            float result = 0;

            //#Add in the contribution to entropy of each action
            foreach (KeyValuePair<string, int> actionTally in actionTallies)
            {
                float proportion = (float)actionTally.Value / exampleCount;
                //Console.WriteLine("Proportion of {0}: {1}", actionTally.Key, proportion.ToString());
                result -= proportion * (float)Math.Log(proportion, 2);
            }

            //# Return the total entropy
            return result;
        }

        public static float entropyOfSets(Dictionary<string, List<Example>> sets, int exampleCount)
        {
            //# Start with zero entropy
            float result = 0;

            //# Get the entropy contribution of each set
            foreach (KeyValuePair<string, List<Example>> set in sets)
            {
                //# Calculate the proportion of the whole in this set
                float proportion = (float)set.Value.Count / exampleCount;
                //# Calculate the entropy contribution
                result -= proportion * entropy(set.Value);
            }

            //#Return the total entropy
            return result;
        }

        public static void makeTree(List<Example> examples, List<string> attributes, DecisionNode decisionNode)
        {
            //# Calculate our initial entropy
            float initialEntropy = entropy(examples);

            //# If we have no entropy, we can’t divide further
            if (initialEntropy <= 0)
            {
                decisionNode.IsLeaf = true;
                int value1Count = 0;
                int value2Count = 0;
                foreach (Example example in examples)
                {
                    if (example.Action.Equals(Example.ActionCollection[0]))
                    {
                        value1Count++;
                    }
                    else
                    {
                        value2Count++;
                    }
                }
                decisionNode.Value = (value1Count > value2Count) ? Example.ActionCollection[0] : Example.ActionCollection[1];
                return;
            }

            //# Find the number of examples
            int exampleCount = examples.Count;

            if (attributes.Count != 0)
            {
                //# Hold the best found split so far
                float bestInformationGain = float.MinValue;
                string bestSplitAttribute = "";
                Dictionary<string, List<Example>> bestSets = new Dictionary<string, List<Example>>();

                //# Go through each attribute
                foreach (string attribute in attributes)
                {
                    //# Perform the split
                    Dictionary<string, List<Example>> sets = splitByAttribute(examples, attribute);
                    //# Find overall entropy and information gain
                    float overallEntropy = entropyOfSets(sets, exampleCount);

                    float informationGain = initialEntropy + overallEntropy;
                    //# Check if we’ve got the best so far
                    if (informationGain > bestInformationGain)
                    {
                        bestInformationGain = informationGain;
                        bestSplitAttribute = attribute;
                        bestSets = sets;
                    }
                }
                //# Set the decision node’s test
                decisionNode.TestValue = bestSplitAttribute;
                //# The list of attributes to pass on down the tree should
                //# have the one we’re using removed
                List<string> newAttributes = attributes.ToList();
                newAttributes.Remove(bestSplitAttribute);

                //# Fill the daughter nodes
                foreach (KeyValuePair<string, List<Example>> set in bestSets)
                {
                    if (set.Value.Count != 0)
                    {
                        //# Find the value for the attribute in this set
                        string attributeValue = set.Value[0].getValue(bestSplitAttribute);
                        //# Create a daughter node for the tree
                        DecisionNode child = new DecisionNode(decisionNode.NodeLevel + 1);
                        //# Add it to the tree
                        decisionNode.setChildNodes(attributeValue, child);
                        //# Recurse the algorithm
                        makeTree(set.Value, newAttributes, child);
                    }
                }
            }
            else
            {
                decisionNode.IsLeaf = true;
                int value1Count = 0;
                int value2Count = 0;
                foreach (Example example in examples)
                {
                    if (example.Action.Equals(Example.ActionCollection[0]))
                    {
                        value1Count++;
                    }
                    else
                    {
                        value2Count++;
                    }
                }
                decisionNode.Value = (value1Count > value2Count) ? Example.ActionCollection[0] : Example.ActionCollection[1];
                int a = 0;
            }
        }

        public static void printTree(DecisionNode node)
        {
            //for (int i = 0; i < node.NodeLevel; i++)
            //{
            //    Console.Write(" ");
            //}
            if (node.IsLeaf)
            {
                Console.ForegroundColor = ConsoleColor.Red;
                Console.WriteLine(" [{0}]", node.Value);
                Console.ResetColor();
            }
            else
            {
                Console.ForegroundColor = ConsoleColor.Green;
                Console.WriteLine("<{0}>", node.TestValue);
                Console.ResetColor();
                foreach (KeyValuePair<string, DecisionNode> child in node.ChildNodes)
                {
                    for (int i = 0; i < child.Value.NodeLevel; i++)
                    {
                        Console.Write(" ");
                    }
                    //Console.ForegroundColor = ConsoleColor.Yellow;
                    Console.Write("+");
                    Console.ForegroundColor = ConsoleColor.Yellow;
                    Console.Write("{0}",child.Key);
                    Console.ResetColor();
                    printTree(child.Value);
                }
            }
        }

        public static bool testNode(Example example,DecisionNode node)
        {
            string testValue = node.TestValue;
            DecisionNode childNode;
            node.ChildNodes.TryGetValue(example.getValue(testValue), out childNode);
            if (childNode == null)
            {
                Console.ForegroundColor = ConsoleColor.Green;
                Console.WriteLine("The {0}'s information is not contained in the tree\n", example.Name);
                Console.ResetColor();
                return true;
            }
            else
            {
                if (childNode.IsLeaf)
                {
                    return example.Action.Equals(childNode.Value);
                }
                else
                {
                    return testNode(example, childNode);
                }
            }
        }
    }
}
